"""
Date: 2022-01-04 16:14:07
LastEditTime: 2022-01-11 15:01:35
Description: voc 格式转换成 coco格式，并且按比例分配到train和val
            所有文件在unlabel文件夹里面用labelimg标注成voc格式

"""

import shutil
import time
import random
import xml.etree.ElementTree as ET
import os
import json
from loguru import logger

annotation_id = 0
category_item_id = -1
image_id = -1

def addImgItem(coco, image_set, file_name, size):
    global image_id
    if file_name is None:
        raise Exception("Could not find filename tag in xml file.")
    if size["width"] is None:
        raise Exception("Could not find width tag in xml file.")
    if size["height"] is None:
        raise Exception("Could not find height tag in xml file.")
    image_id += 1
    image_item = dict()
    image_item["id"] = image_id
    image_item["file_name"] = file_name
    image_item["width"] = size["width"]
    image_item["height"] = size["height"]
    coco["images"].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(coco, _id, category_id, bbox):
    global annotation_id
    global image_id
    image_id = _id

    annotation_item = dict()
    annotation_item["segmentation"] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])

    annotation_item["segmentation"].append(seg)

    annotation_item["area"] = bbox[2] * bbox[3]
    annotation_item["iscrowd"] = 0
    annotation_item["ignore"] = 0
    annotation_item["image_id"] = image_id
    annotation_item["bbox"] = bbox
    annotation_item["category_id"] = category_id

    annotation_id += 1

    annotation_item["id"] = annotation_id
    coco["annotations"].append(annotation_item)

def parseXmlFiles(xml_path, json_file, _id, COCO_CLASSES):
    global image_id
    image_id = _id

    coco = dict()
    coco["images"] = []
    coco["annotations"] = []
    coco["categories"] = []

    category_set = dict()
    image_set = set()

    for idx, item in enumerate(COCO_CLASSES):
        category_set[item] = idx
        # 添加到categories中,保证categories里的顺序不变
        category_item = dict()
        category_item["supercategory"] = "none"
        category_item["id"] = idx
        category_item["name"] = item
        coco["categories"].append(category_item)

    file_all = os.listdir(xml_path)

    xml_file_counter=0

    for f in file_all:
        if not f.endswith(".xml"):
            continue
        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size["width"] = None
        size["height"] = None
        size["depth"] = None

        xml_file = os.path.join(xml_path, f)
        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != "annotation":
            raise Exception(
                "pascal voc xml root element should be annotation, rather than {}".format(
                    root.tag
                )
            )

        xml_file_counter+=1

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == "folder":
                continue

            if elem.tag == "filename":  # 不使用tag filename 而是使用和.xml同名的.jpg文件名
                # file_name = elem.text
                name, ext = os.path.splitext(f)  # 分离文件名和扩展名
                file_name = name + ".jpg"
                if file_name not in file_all:    # 判断是否已经存在相应的jpg文件(只支持jpg格式)
                    logger.error("file ", file_name, " not exist")
                    continue

            # add img item only after parse <size> tag
            elif (
                current_image_id is None
                and file_name is not None
                and size["width"] is not None
            ):
                if file_name not in image_set:
                    current_image_id = addImgItem(coco, image_set, file_name, size)

                    # print("add image with {} and {}".format(file_name, size))

                else:
                    raise Exception("duplicated image: {}".format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>

            for subelem in elem:
                bndbox["xmin"] = None
                bndbox["xmax"] = None
                bndbox["ymin"] = None
                bndbox["ymax"] = None

                current_sub = subelem.tag
                if current_parent == "object" and subelem.tag == "name":
                    object_name = subelem.text

                    # assert object_name in category_set, "出现未在coco_class中的类别名"
                    # if object_name not in category_set:
                    # current_category_id = addCatItem(coco, category_set, object_name) #新category 添加

                    if object_name not in category_set:
                        print("类别‘%s’不在coco_class中，跳过" % (object_name))
                        continue
                    current_category_id = category_set[object_name]

                elif current_parent == "size":
                    if size[subelem.tag] is not None:
                        raise Exception("xml structure broken at size tag.")
                    size[subelem.tag] = int(subelem.text)
                    continue

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == "bndbox":
                        if bndbox[option.tag] is not None:
                            raise Exception("xml structure corrupted at bndbox tag.")
                        bndbox[option.tag] = int(option.text)

                # only after parse the <object> tag
                if bndbox["xmin"] is not None:
                    # if object_name is None:
                    #     raise Exception("xml structure broken at bndbox tag")
                    # if current_image_id is None:
                    #     raise Exception("xml structure broken at bndbox tag")
                    # if current_category_id is None:
                    #     raise Exception("xml structure broken at bndbox tag")

                    if object_name is None or current_image_id is None or current_category_id is None:
                        print("xml structure broken at bndbox tag")
                        continue

                    bbox = []
                    # x
                    bbox.append(bndbox["xmin"])
                    # y
                    bbox.append(bndbox["ymin"])
                    # w
                    bbox.append(bndbox["xmax"] - bndbox["xmin"])
                    # h
                    bbox.append(bndbox["ymax"] - bndbox["ymin"])

                    # print(
                    #     "add annotation with {},{},{},{}".format(
                    #         object_name, current_image_id, current_category_id, bbox
                    #     )
                    # )

                    addAnnoItem(
                        coco, current_image_id, current_category_id, bbox
                    )
                    
    logger.info("process {} xml file".format(xml_file_counter))

    json.dump(coco, open(json_file, "w"))


def splitLabeled(val_percent, labeled_dir, train_dir, val_dir):
    maxId = __getMaxFileNameId(train_dir)
    maxId = __getMaxFileNameId(val_dir, maxId)

    print("Max file id=", str(maxId))

    # 创建目录操作函数
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(val_dir):
        os.makedirs(val_dir)

    random.seed(int(time.time()))

    fileAll = os.listdir(labeled_dir)

    for fileName in fileAll:
        if not fileName.endswith(".jpg"):
            continue
        name, ext = os.path.splitext(fileName)  # 分离文件名和扩展名
        xml_name = name + ".xml"  # 判断是否已经存在相应的xml文件
        if xml_name not in fileAll:
            continue

        # 决定要转移到的目标目录
        targetDir = val_dir
        r = random.random()
        if r > val_percent:
            targetDir = train_dir

        # 更改名称移动
        maxId = maxId + 1
        # .jpg
        target_name = str(maxId).zfill(6)
        full_path_jpg = os.path.join(labeled_dir, fileName)  # 将文件名与文件目录连接起来，形成完整路径
        target_path = os.path.join(targetDir, (target_name + ext))

        shutil.move(full_path_jpg, target_path)
        print(full_path_jpg, " to ", target_path)

        # .xml
        full_path = os.path.join(labeled_dir, xml_name)  # 将文件名与文件目录连接起来，形成完整路径
        target_path = os.path.join(targetDir, (target_name + ".xml"))
        shutil.move(full_path, target_path)


# 获取名字最大的一个.jpg文件（用于文件命名，避免重复）
def __getMaxFileNameId(filePath: str = None, fileMaxId: int = 0):
    assert filePath is not None, "file path is None"
    fileList = os.listdir(filePath)
    maxId = fileMaxId
    for fileWithPath in fileList:
        if not fileWithPath.endswith(".jpg"):
            continue
        # dir, fileName = os.path.split(fileWithPath)  # 分离路径和文件名
        name, ext = os.path.splitext(fileWithPath)  # 分离文件名和扩展名
        if not name.isdigit():  # 判断是否是数字
            continue
        n = int(name)
        if n > maxId:
            maxId = n
    # print(filePath, str(maxId))
    return maxId


def delImgCache(path):
    for maindir, subdir, file_name_list in os.walk(path):
        for filename in file_name_list:
            ext = os.path.splitext(filename)[1]
            if '.array' == ext:
                apath = os.path.join(maindir, filename)
                os.unlink(apath)
                print("del image cache file->", apath)

def reduceData(train_dir, val_dir,max_data=500,val_percent=0.1,recycle_dir='d:/ai/yolox/datasets/coco/recycle'):
    maxId = __getMaxFileNameId(train_dir)
    maxId = __getMaxFileNameId(val_dir, maxId)
    os.makedirs(recycle_dir, exist_ok=True)

    train_all=os.listdir(train_dir)
    val_all=os.listdir(val_dir)

    if len(train_all)<1 or len(val_all)<1:
        logger.warning("train of val data is null")
        return

    max_data*=2 #两个文件

    if len(train_all)+len(val_all)<max_data:#低于max，不处理
        logger.info("data less than {}" .format(int(max_data/2)))
        return
    
    #编号最大的100个不处理，其余的按比例转移到recycle文件夹
    train_left=float(max_data)*(1-val_percent)  
    n_move=len(train_all)-train_left
    n_process=n_move+(train_left-100*2*(1-val_percent)) 
    remove_rate=n_move/n_process
    logger.debug("train remove rate : {}".format(remove_rate))     
    __moveJpgXml(train_dir,recycle_dir,remove_rate,maxId)#移动

    val_left=max_data*val_percent
    n_move=len(val_all)-val_left
    n_process=n_move+(val_left-100*2*val_percent) 
    remove_rate=n_move/n_process
    logger.debug("val remove rate : {}".format(remove_rate))    
    __moveJpgXml(val_dir,recycle_dir,remove_rate,maxId) #移动val


def __moveJpgXml(source_dir,target_dir,remove_rate,max_id):
    remove_count=0
    if remove_rate<0.000001:
        return
    
    file_all=os.listdir(source_dir)
    for filename in file_all:
        if not filename.endswith(".jpg"):
            continue
        name, ext = os.path.splitext(filename)  # 分离文件名和扩展名
        xml_name = name + ".xml"  # 判断是否已经存在相应的xml文件
        if xml_name not in file_all:
            continue

        if not name.isdigit():  # 判断是否是数字
            continue
        n = int(name)
        if n > max_id-100: #编号最大的100个文件（最新）不处理
            continue

        r = random.random()
        if r > remove_rate:
            continue

        # .jpg
        full_path_jpg = os.path.join(source_dir, filename)  # 将文件名与文件目录连接起来，形成完整路径
        target_path = os.path.join(target_dir, filename)

        shutil.move(full_path_jpg, target_path)
        logger.info("{} to {}".format(full_path_jpg, target_path))

        remove_count+=1

        # .xml
        name, ext = os.path.splitext(filename)  # 分离文件名和扩展名
        full_path = os.path.join(source_dir, name+".xml")  # 将文件名与文件目录连接起来，形成完整路径
        target_path = os.path.join(target_dir,name+'.xml')
        shutil.move(full_path, target_path)

    logger.info("remove count : {}".format(remove_count))

if __name__ == "__main__":

    # 定义类别，要与yolox/data/datasets/coco_classes.py一致
    COCO_CLASSES = (
        "slipper",
        "outsole",
        "face",
        "pill"
    )

    new_labeled_img_dir = "d:/ai/yolox/datasets/coco/unlabel" # 新标注的文件
    train_dir = "d:/ai/yolox/datasets/coco/train2017"         # 目标文件夹train
    val_dir = "d:/ai/yolox/datasets/coco/val2017"             # 目标文件夹val
    val_percent = 0.1                                         # train val 比例
    cache_dir = 'd:/ai/yolox/datasets/coco'                   # 缓存图片的文件路径！！！！！一定要删除，被坑一个星期！！！！！！！！

    # 按比例随机拆分到到目标目录，会同时移动.jpg和.xml文件
    splitLabeled(val_percent, new_labeled_img_dir, train_dir, val_dir)

    #回收旧数据,提高新数据占比，最新100个数据不参与回收
    reduceData(train_dir,val_dir,max_data=500,val_percent=val_percent,recycle_dir='d:/ai/yolox/datasets/coco/recycle')

    # 生成json
    # train
    json_file = "d:/ai/yolox/datasets/coco/annotations/instances_train2017.json"
    image_id = 202200000
    parseXmlFiles(train_dir, json_file, image_id, COCO_CLASSES)

    # eval
    json_file = "d:/ai/yolox/datasets/coco/annotations/instances_val2017.json"
    image_id = 0
    parseXmlFiles(val_dir, json_file, image_id, COCO_CLASSES)

    # 删除 cache 文件
    delImgCache(cache_dir)
